{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import argparse\n",
    "import random\n",
    "import torch.nn as nn\n",
    "import torch.nn.parallel\n",
    "import torch.nn.functional as F\n",
    "import torch.backends.cudnn as cudnn\n",
    "import torch.optim as optim\n",
    "from torchvision import transforms, utils\n",
    "from torch.autograd import Variable\n",
    "from einops import rearrange, reduce, repeat"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "parser = argparse.ArgumentParser()\n",
    "parser.add_argument('--dataroot', required=True, help='path to dataset')\n",
    "parser.add_argument('--workers', type=int, help='number of data loading workers', default=8)\n",
    "parser.add_argument('--batchSize', type=int, default=64, help='input batch size')\n",
    "parser.add_argument('--niter', type=int, default=200, help='number of epochs to train for')\n",
    "parser.add_argument('--lr', type=float, default=0.001, help='learning rate')\n",
    "parser.add_argument('--beta1', type=float, default=0.9, help='beta1 for adam. ')\n",
    "parser.add_argument('--cuda', action='store_true', help='enables cuda')\n",
    "parser.add_argument('--ngpu', type=int, default=2, help='number of GPUs to use')\n",
    "parser.add_argument('--manualSeed', type=int, help='manual seed')\n",
    "opt = parser.parse_args(['--dataroot','./data','--cuda'])\n",
    "#print(opt)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "os.environ[\"CUDA_VISIBLE_DEVICES\"] = \"0,1\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if opt.manualSeed is None:\n",
    "    opt.manualSeed = random.randint(1, 10000)\n",
    "print(\"Random Seed: \", opt.manualSeed)\n",
    "random.seed(opt.manualSeed)\n",
    "torch.manual_seed(opt.manualSeed)\n",
    "if opt.cuda:\n",
    "    torch.cuda.manual_seed_all(opt.manualSeed)\n",
    "cudnn.benchmark = True\n",
    "if torch.cuda.is_available() and not opt.cuda:\n",
    "    print(\"WARNING: You have a CUDA device, so you should probably run with --cuda\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class AdditiveAttention(nn.Module):\n",
    "    def __init__(self, key_size, query_size, num_hiddens):\n",
    "        super(AdditiveAttention, self).__init__()\n",
    "        self.w_k = nn.Linear(key_size, num_hiddens, bias=False)\n",
    "        self.w_q = nn.Linear(query_size, num_hiddens, bias=False)\n",
    "        self.w_v = nn.Linear(num_hiddens, 1, bias=False)\n",
    "        self.softmax = nn.Softmax(dim = 1)\n",
    "\n",
    "    def forward(self, queries, keys):\n",
    "        # queries's shape: [bs, query_size]\n",
    "        # keys's shape: [bs, seq_len, query_size]\n",
    "        queries = queries.unsqueeze(1)# [bs, query_size]->[bs, 1, query_size]\n",
    "        queries = self.w_q(queries)\n",
    "        keys = self.w_k(keys)\n",
    "        features = torch.tanh(queries + keys)\n",
    "        scores = self.w_v(features)\n",
    "        weights = self.softmax(scores)\n",
    "\n",
    "        return weights\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cfg = {\n",
    "    'A': [32, 32, 'M',64, 64, 'M']\n",
    "}\n",
    "\n",
    "def make_layers(cfg, in_channels = 3, batch_norm=True):\n",
    "    layers = [] \n",
    "    for v in cfg:\n",
    "        if v == 'M':\n",
    "            layers += [nn.MaxPool1d(kernel_size=2, stride=2)]\n",
    "        else:\n",
    "            conv1d = nn.Conv1d(in_channels, v, kernel_size=3, stride=1, padding=1,dilation=1)\n",
    "            if batch_norm:\n",
    "                layers += [conv1d, nn.BatchNorm1d(v), nn.ReLU(inplace=True)]\n",
    "            else:\n",
    "                layers += [conv1d, nn.ReLU(inplace=True)]\n",
    "            in_channels = v\n",
    "    return nn.Sequential(*layers)\n",
    "\n",
    "def conv3x3(in_planes, out_planes, stride=1, groups=1, dilation=1):\n",
    "    \"\"\"3x3 convolution with padding\"\"\"\n",
    "    return nn.Conv1d(in_planes, out_planes, kernel_size=3, stride=stride,\n",
    "                     padding=dilation, groups=groups, bias=False, dilation=dilation)\n",
    "\n",
    "def conv1x1(in_planes, out_planes, stride=1):\n",
    "    \"\"\"1x1 convolution\"\"\"\n",
    "    return nn.Conv1d(in_planes, out_planes, kernel_size=1, stride=stride, bias=False)\n",
    "\n",
    "class BasicBlock(nn.Module):\n",
    "    def __init__(self, inplanes, planes, shortcut = True):\n",
    "        super(BasicBlock, self).__init__()\n",
    "        self.downsample = None\n",
    "\n",
    "        self.conv1 = conv3x3(inplanes, planes)\n",
    "        self.bn1 = nn.BatchNorm1d(planes)\n",
    "        self.relu = nn.ReLU(inplace=True)\n",
    "        self.conv2 = conv3x3(planes, planes)\n",
    "        self.bn2 = nn.BatchNorm1d(planes)\n",
    "        self.conv3 = conv3x3(planes, planes)\n",
    "        self.bn3 = nn.BatchNorm1d(planes)\n",
    "        self.shortcut = shortcut\n",
    "        if inplanes != planes:\n",
    "            self.downsample = nn.Sequential(\n",
    "                conv1x1(inplanes, planes),\n",
    "                nn.BatchNorm1d(planes),\n",
    "            )\n",
    "\n",
    "    def forward(self, x):\n",
    "        identity = x #\n",
    "\n",
    "        out = self.conv1(x)\n",
    "        out = self.bn1(out)\n",
    "        out = self.relu(out)\n",
    "\n",
    "        out = self.conv2(out)\n",
    "        out = self.bn2(out)\n",
    "        out = self.relu(out)\n",
    "\n",
    "        out = self.conv3(out)\n",
    "        out = self.bn3(out)\n",
    "        \n",
    "        if self.downsample is not None:\n",
    "            identity = self.downsample(x)\n",
    "                    \n",
    "        if self.shortcut:\n",
    "            out += identity\n",
    "\n",
    "        out = self.relu(out)\n",
    "        return out\n",
    "\n",
    "class CNN_Net(nn.Module):\n",
    "    def __init__(self, inplanes=3):\n",
    "        super(CNN_Net, self).__init__()\n",
    "        self.features = nn.Sequential(\n",
    "            make_layers(cfg['A'], in_channels = inplanes, batch_norm=True),\n",
    "            BasicBlock(inplanes=64, planes=128),\n",
    "            nn.MaxPool1d(kernel_size=2, stride=3),\n",
    "            BasicBlock(inplanes=128, planes=256),\n",
    "            nn.MaxPool1d(kernel_size=2, stride=3),\n",
    "            BasicBlock(inplanes=256, planes=256),\n",
    "            nn.MaxPool1d(kernel_size=2, stride=3)\n",
    "        )\n",
    "\n",
    "    def forward(self, x):\n",
    "        x = self.features(x)\n",
    "        return x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class PositionalEncoding(nn.Module):\n",
    "    def __init__(self, d_model, dropout=0.1, max_len=5000):\n",
    "        super(PositionalEncoding, self).__init__()\n",
    "        \n",
    "        self.dropout = nn.Dropout(p=dropout)\n",
    "        pe = torch.zeros(max_len, d_model)\n",
    "        position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1)\n",
    "        div_term = torch.exp(torch.arange(0, d_model, 2).float() * (-math.log(10000.0) / d_model))\n",
    "        pe[:, 0::2] = torch.sin(position * div_term)\n",
    "      \n",
    "        pe[:, 1::2] = torch.cos(position * div_term)\n",
    "        \n",
    "        pe = pe.unsqueeze(0)\n",
    "        self.register_buffer('pe', pe)  \n",
    "        \n",
    "    def forward(self, x):\n",
    "        \"\"\"\n",
    "        x: [batch_size, seq_len, d_model]\n",
    "        \"\"\"\n",
    "        #x = x + self.pe[:,:x.size(1)]\n",
    "        x = x + Variable(self.pe[:, :x.size(1)], requires_grad=False)\n",
    "        return self.dropout(x)\n",
    "\n",
    "class MultiHeadAttention(nn.Module):\n",
    "    def __init__(self, emb_size, num_heads, dropout):\n",
    "        super().__init__()\n",
    "        self.emb_size = emb_size\n",
    "        self.num_heads = num_heads\n",
    "        self.keys = nn.Linear(emb_size, emb_size)\n",
    "        self.queries = nn.Linear(emb_size, emb_size)\n",
    "        self.values = nn.Linear(emb_size, emb_size)\n",
    "        self.att_drop = nn.Dropout(dropout)\n",
    "        self.projection = nn.Linear(emb_size, emb_size)\n",
    "\n",
    "    def forward(self, x, mask= None):\n",
    "        queries = rearrange(self.queries(x), \"b n (h d) -> b h n d\", h=self.num_heads)\n",
    "        keys = rearrange(self.keys(x), \"b n (h d) -> b h n d\", h=self.num_heads)\n",
    "        values = rearrange(self.values(x), \"b n (h d) -> b h n d\", h=self.num_heads)\n",
    "        energy = torch.einsum('bhqd, bhkd -> bhqk', queries, keys)  \n",
    "        if mask is not None:\n",
    "            fill_value = torch.finfo(torch.float32).min\n",
    "            energy.mask_fill(~mask, fill_value)\n",
    "\n",
    "        scaling = self.emb_size ** (1 / 2)\n",
    "        att = F.softmax(energy / scaling, dim=-1)\n",
    "        att = self.att_drop(att)\n",
    "        out = torch.einsum('bhal, bhlv -> bhav ', att, values)\n",
    "        out = rearrange(out, \"b h n d -> b n (h d)\")\n",
    "        out = self.projection(out)\n",
    "        return out\n",
    "\n",
    "\n",
    "class ResidualAdd(nn.Module):\n",
    "    def __init__(self, fn):\n",
    "        super().__init__()\n",
    "        self.fn = fn\n",
    "\n",
    "    def forward(self, x, **kwargs):\n",
    "        res = x\n",
    "        x = self.fn(x, **kwargs)\n",
    "        x += res\n",
    "        return x\n",
    "\n",
    "\n",
    "class FeedForwardBlock(nn.Sequential):\n",
    "    def __init__(self, emb_size, expansion, drop_p):\n",
    "        super().__init__(\n",
    "            nn.Linear(emb_size, expansion * emb_size),\n",
    "            nn.ReLU(),\n",
    "            #nn.GELU(),\n",
    "            nn.Dropout(drop_p),\n",
    "            nn.Linear(expansion * emb_size, emb_size),\n",
    "        )\n",
    "\n",
    "\n",
    "class GELU(nn.Module):\n",
    "    def forward(self, input):\n",
    "        return input*0.5*(1.0+torch.erf(input/math.sqrt(2.0)))\n",
    "\n",
    "class TransformerEncoderBlock(nn.Sequential):\n",
    "    def __init__(self,\n",
    "                 emb_size,\n",
    "                 num_heads=4,\n",
    "                 drop_p=0.1,\n",
    "                 forward_expansion=2,\n",
    "                 forward_drop_p=0.1):\n",
    "        super().__init__(\n",
    "            ResidualAdd(nn.Sequential(\n",
    "                nn.LayerNorm(emb_size),\n",
    "                MultiHeadAttention(emb_size, num_heads, drop_p),\n",
    "                nn.Dropout(drop_p)\n",
    "            )),\n",
    "            ResidualAdd(nn.Sequential(\n",
    "                nn.LayerNorm(emb_size),\n",
    "                FeedForwardBlock(\n",
    "                    emb_size, expansion=forward_expansion, drop_p=forward_drop_p),\n",
    "                nn.Dropout(drop_p)\n",
    "            )\n",
    "            ))\n",
    "\n",
    "class TransformerEncoder(nn.Module):\n",
    "    def __init__(self, depth, emb_size, num_heads, pos=True):\n",
    "        super(TransformerEncoder, self).__init__()\n",
    "        self.emb = PositionalEncoding(d_model=emb_size,dropout=0.1, max_len=256)\n",
    "        self.layers = nn.ModuleList([TransformerEncoderBlock(emb_size, num_heads = num_heads) for _ in range(depth)])\n",
    "        self.pos = pos\n",
    "\n",
    "    def forward(self, x):\n",
    "        \n",
    "        if self.pos:\n",
    "            x = self.emb(x)\n",
    "        \n",
    "        for layer in self.layers:\n",
    "            x = layer(x)\n",
    "        \n",
    "        return x\n",
    "\n",
    "rnn = TransformerEncoder(depth=3, emb_size=256, num_heads = 4, pos=True)\n",
    "x = torch.randn(5,111,256)\n",
    "out = rnn(x)\n",
    "print(out.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class HGCTNet(nn.Module):\n",
    "    def __init__(self, num_classes=2, num_hid = 8, depth = 3, num_heads = 4, hga = True, trans = True, init_weights = True):\n",
    "        super(HGCTNet, self).__init__()\n",
    "        self.hga = hga\n",
    "        self.trans = trans\n",
    "        self.cnn = CNN_Net()\n",
    "        self.regressor = nn.Sequential(\n",
    "            nn.Linear(256, 128),\n",
    "            #nn.Linear(256+51, 128),\n",
    "            nn.ReLU(True),\n",
    "            nn.Dropout(),\n",
    "            nn.Linear(128, num_classes)\n",
    "        )\n",
    "        \n",
    "        # AdditiveAttention\n",
    "        self.cnn_attention = AdditiveAttention(key_size = 111, query_size = 47, num_hiddens = num_hid)\n",
    "        \n",
    "        self.avgpool = nn.AdaptiveAvgPool1d(1)\n",
    "        self.conv1x1 = conv1x1(in_planes=11, out_planes=1) \n",
    "        self.emblay = nn.Sequential(\n",
    "            nn.Linear(51,128, bias=True),\n",
    "            nn.ReLU(True),\n",
    "            nn.Dropout(0.5),\n",
    "            nn.Linear(128, 256, bias=True),\n",
    "        )\n",
    "        self.rnn = TransformerEncoder(depth=depth, emb_size=256, num_heads = num_heads, pos=True)\n",
    "        self.balance_alpha = torch.nn.Parameter(torch.randn(2))\n",
    "        \n",
    "        if init_weights:\n",
    "            self._initialize_weights()\n",
    "\n",
    "    def forward(self, x, feas):\n",
    "        basebp = feas[:,0:2]\n",
    "        demos = feas[:,2:6]\n",
    "        hand_feas = feas[:,6:53]\n",
    "\n",
    "        # CNN module\n",
    "        x_s = self.cnn(x)\n",
    "        \n",
    "        # HGA module\n",
    "        if self.hga: \n",
    "            alpha = self.cnn_attention(hand_feas, x_s)\n",
    "            x_s = torch.mul(alpha, x_s)\n",
    "        x_res = self.avgpool(x_s)\n",
    "        x_res = x_res.view(x_res.size(0), -1)\n",
    "        \n",
    "        # Transformer module\n",
    "        if self.trans:\n",
    "            x_s = x_s.permute(0, 2, 1)\n",
    "            x_s = self.rnn(x_s)\n",
    "            x_s = x_s.permute(0, 2, 1)\n",
    "            x_s = self.avgpool(x_s)\n",
    "            x_s = x_s[:,:,0]\n",
    "        else:\n",
    "            x_s = self.avgpool(x_s)\n",
    "            x_s = x_s[:,:,0]\n",
    "        \n",
    "        # Feature fusion module\n",
    "        x = x_s + x_res\n",
    "        x = torch.mul(x,self.emblay(torch.cat([demos,hand_feas],dim=1)))\n",
    "        #x = torch.cat([x,torch.cat([demos,hand_feas],dim=1)],dim=1)\n",
    "        \n",
    "        # regression\n",
    "        x = self.regressor(x)\n",
    "        #x = x*0.5 + basebp*(1-0.5)\n",
    "        x = x*self.balance_alpha + basebp*(1-self.balance_alpha)\n",
    "\n",
    "        return x\n",
    "   \n",
    "    def _initialize_weights(self):\n",
    "        for m in self.modules():\n",
    "            if isinstance(m, nn.Conv1d): \n",
    "                nn.init.kaiming_normal_(m.weight.data,mode='fan_out')\n",
    "                if m.bias is not None:\n",
    "                    m.bias.data.zero_()\n",
    "            elif isinstance(m, nn.BatchNorm1d):\n",
    "                m.weight.data.fill_(1)\n",
    "                m.bias.data.zero_()\n",
    "            elif isinstance(m, nn.Linear):\n",
    "                nn.init.kaiming_normal_(m.weight.data,mode='fan_out')\n",
    "                if m.bias is not None:\n",
    "                    m.bias.data.zero_()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# example of blood pressure estimation in CAS-BP dataset\n",
    "# input: \n",
    "#      sigs: [bs, channel, dimension]->[bs, 3, 12000], input three signals (ecg+ppg+ppw), each 60 seconds in length (sampling rate 200Hz)\n",
    "#      feas: [bs, dimension]->[bs, 53] includes: {baseline sbp, baseline dbp, demos (age, gender, BMI, history of hypertension), \n",
    "#                                                   47 handcrafted features calculated from the three signals}\n",
    "#\n",
    "# output: [bs, dimension]->[bs, 53] estimated SBP and DBP\n",
    "\n",
    "model = HGCTNet(num_classes=2, num_hid = 16, depth = 2, num_heads = 2, hga = True, trans = True)\n",
    "model.cuda()\n",
    "sigs = Variable(torch.randn(5,3,12000)).cuda()\n",
    "feas = Variable(torch.randn(5,53)).cuda()\n",
    "estbp = model(sigs, feas)\n",
    "#print(estbp)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3.10.9 ('pytorch')",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.9"
  },
  "orig_nbformat": 4,
  "vscode": {
   "interpreter": {
    "hash": "45c921d51f671e7c0f576ae058dd5aa62c17caf3109d6bd9b696251b4cd980ad"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
